import cv2
from ultralytics import YOLO


def pre_video():

    model_path = r'C:\Users\64531\Desktop\03smallPaper_all\0001源网络Ciou训练有裁剪的数据\weights\best.pt'  # Change this to your YOLOv8 model's path
    video_path = r"C:\Users\64531\Desktop\weed_vedio.mp4"  # Change this to your video's path

    # Load the trained YOLOv8 model
    model = YOLO(model_path)
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        print("Error: Could not open video.")
        exit()


    # Process video frames
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            print("Finished processing video.")
            break

        if ret:

            results = model(frame)

            annotated_frame = results[0].plot()

            cv2.imshow("YOLOv8 Inference", annotated_frame)

            # EXC退出
            if cv2.waitKey(1) == 27:
                break

        else:
            break

    cap.release()
    cv2.destroyAllWindows()


def pre_camera():

    model_path = r'C:\Users\64531\Desktop\03smallPaper_all\0001源网络Ciou训练有裁剪的数据\weights\best.pt'  # Change this to your YOLOv8 model's path
    model_path = r'D:\pythonProject\ultralytics4.12\yolov8n.pt'  # Change this to your YOLOv8 model's path

    camera_no = 0

    # Load the trained YOLOv8 model
    model = YOLO(model_path)

    cap = cv2.VideoCapture(camera_no)

    if not cap.isOpened():
        print("Error: Could not open video.")
        exit()


    # Process video frames
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            print("Finished processing video.")
            break

        if ret:

            results = model(frame)

            annotated_frame = results[0].plot()

            cv2.imshow("YOLOv8 Inference", annotated_frame)

            # EXC退出
            if cv2.waitKey(1) == 27:
                break

        else:
            break

    cap.release()
    cv2.destroyAllWindows()




if __name__ == '__main__':
    # pre_video()
    pre_camera()
